from flask import Flask, request, jsonify, send_file
import requests
import os
from io import BytesIO
from PIL import Image
import base64

app = Flask(__name__)

# 使用免费的AI图片生成API（示例）
API_URL = "https://api.openai.com/v1/images/generations"
API_KEY = "your-api-key-here"  # 用户需要替换为自己的API密钥

@app.route('/')
def index():
    return send_file('index.html')

@app.route('/generate', methods=['POST'])
def generate_image():
    try:
        data = request.json
        prompt = data.get('prompt', '')
        
        if not prompt:
            return jsonify({'error': '请输入图片描述'}), 400
        
        # 这里使用一个简单的占位符图片生成逻辑
        # 实际应用中应该调用真实的AI API
        placeholder_url = f"https://via.placeholder.com/512x512/4CAF50/ffffff?text={prompt.replace(' ', '+')}"
        
        # 下载占位符图片
        response = requests.get(placeholder_url)
        img = Image.open(BytesIO(response.content))
        
        # 保存图片
        output_path = 'generated_image.png'
        img.save(output_path)
        
        # 转换为base64用于前端显示
        buffered = BytesIO()
        img.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()
        
        return jsonify({
            'success': True,
            'image_data': f'data:image/png;base64,{img_str}',
            'filename': output_path
        })
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/download/<filename>')
def download_image(filename):
    try:
        return send_file(filename, as_attachment=True)
    except Exception as e:
        return jsonify({'error': str(e)}), 404

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)
